#include "Opt.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/IntrinsicInst.h"

using namespace llvm;

// should read from file
OptForAsan::OptForAsan(llvm::DenseMap<const llvm::AllocaInst *, bool> &ProcessedAllocas): processAllocas(ProcessedAllocas) {
    processAllocas.clear();
}

OptForAsan::~OptForAsan() {
    //processAllocas.clear();
}

std::unordered_set<const llvm::Instruction *> OptForAsan::ApplyOnFunc(llvm::Function &F, const llvm::TargetLibraryInfo *TLI) {
    // errs() << "OptForAsan::ApplyOnFunc\n";
    _to_instrument.clear();
    bool need = DoFunctionExcludeOpt(F);
    if (need) {
        DoStackSafeVariableOpt(F);
        DoHeapSafeVariableOpt(F);
    }

    return _to_instrument;
}

bool OptForAsan::DoFunctionExcludeOpt(llvm::Function &F) {
    std::string name = F.getName();
    if (name.find("main") != std::string::npos && name != "main") {
        F.addFnAttr(llvm::Attribute::SanitizeAddress);
        return true;
    }
    return false;
}

void OptForAsan::DoStackSafeVariableOpt(llvm::Function &F) {
    for (Instruction &I : F.getEntryBlock()) {
        IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
        if (II != NULL) {
            
            for (Value *Arg : II->arg_operands()) {
                AllocaInst *AI = dyn_cast<AllocaInst>(Arg->stripPointerCasts());
                // assert(AI && AI->isStaticAlloca() && "non-static alloca arg to localescape");
                if (AI != NULL) {
                    //errs() << I << AI << "\n";
                    //processAllocas[AI] = false;
                }
            }
        }
    }
}

void OptForAsan::DoHeapSafeVariableOpt(llvm::Function &F) {
    // pass 1
    checkReleasedVar(F);

    

    errs() << "digraph " + F.getName() + "{\n";
    errs() << "\n";
    for (auto block = F.getBasicBlockList().begin(); block != F.getBasicBlockList().end(); block++) {
        for (auto inst = block->begin(); inst != block->end(); inst++) {
            if(isa<CallInst>(inst) || isa<InvokeInst>(inst)) {
               errs() << "Is function call: " << *inst << "\n";
            }
            for (Use &U:inst->operands()) {
                Value *v = U.get();
                if (dyn_cast<Instruction>(v)) {
                    errs() << "\"" << *dyn_cast<Instruction>(v) << "\"" << " -> " << "\"" << *inst << "\"" << ";\n";
                }
                if (v->getName() != "") {
                    errs() << "\"" << v->getName() << "\"" << " -> " << "\"" << *inst << "\"" << ";\n";
                    errs() << "\"" << v->getName() << "\"" << " [ color = red ]\n";
                }
            }

           
        }
    }
    errs() << "\n}\n";
}